import threading
import time
import json
import os
import uuid
from datetime import datetime
from utils.logger import get_logger
from config import Config
from database import TaskDAO
from video_processor import VideoProcessor

logger = get_logger(__name__)


class TaskManager:
    """任务管理器"""

    def __init__(self, task_dao: TaskDAO):
        self.task_dao = task_dao
        self.video_processor = VideoProcessor()
        self.tasks = {}  # 内存中的任务状态
        self.task_lock = threading.Lock()

    def create_task(self, human_file_id, background_file_id, params, output_name):
        """创建新任务"""
        task_id = str(uuid.uuid4())

        # 保存到内存
        with self.task_lock:
            self.tasks[task_id] = {
                'task_id': task_id,
                'status': 'pending',
                'progress': 0,
                'message': '任务已创建',
                'created_at': datetime.now().isoformat(),
                'updated_at': datetime.now().isoformat(),
                'human_file_id': human_file_id,
                'background_file_id': background_file_id,
                'params': params,
                'output_name': output_name
            }

        # 保存到数据库
        success = self.task_dao.create_task(
            task_id, human_file_id, background_file_id,
            json.dumps(params), output_name, 'pending'
        )

        if not success:
            # 如果数据库保存失败，从内存中移除
            with self.task_lock:
                self.tasks.pop(task_id, None)
            return None

        return task_id

    def get_task_status(self, task_id):
        """获取任务状态"""
        with self.task_lock:
            return self.tasks.get(task_id)

    def update_task_progress(self, task_id, progress, message, status='processing'):
        """更新任务进度"""
        # 更新内存中的状态
        with self.task_lock:
            if task_id in self.tasks:
                self.tasks[task_id].update({
                    'progress': progress,
                    'message': message,
                    'status': status,
                    'updated_at': datetime.now().isoformat()
                })

        # 更新数据库
        self.task_dao.update_task_status(task_id, status, progress)

    def complete_task(self, task_id, result_path, preview_path, processing_time):
        """完成任务"""
        # 更新内存状态
        with self.task_lock:
            if task_id in self.tasks:
                self.tasks[task_id].update({
                    'status': 'completed',
                    'progress': 100,
                    'message': '处理完成',
                    'result_path': result_path,
                    'preview_path': preview_path,
                    'result_url': f"/api/download/{os.path.basename(result_path)}",
                    'updated_at': datetime.now().isoformat()
                })

        # 更新数据库
        self.task_dao.complete_task(task_id, result_path, preview_path, processing_time)

    def fail_task(self, task_id, error_message, processing_time=None):
        """任务失败"""
        # 更新内存状态
        with self.task_lock:
            if task_id in self.tasks:
                self.tasks[task_id].update({
                    'status': 'failed',
                    'message': error_message,
                    'error': error_message,
                    'updated_at': datetime.now().isoformat()
                })

        # 更新数据库
        self.task_dao.update_task_status(task_id, 'failed', error_message=error_message)
        if processing_time:
            # 更新处理时间
            with self.task_dao.db.get_connection() as conn:
                conn.execute('''
                    UPDATE tasks SET processing_time_seconds = ? WHERE task_id = ?
                ''', (processing_time, task_id))
                conn.commit()

    def start_processing(self, task_id, human_video_path, background_path, output_path, params):
        """开始处理任务"""
        # 检查文件是否为视频背景
        is_video_bg = background_path.lower().endswith(('.mp4', '.mov', '.avi', '.mkv'))

        # 启动处理线程
        thread = threading.Thread(
            target=self._process_video_task,
            args=(task_id, human_video_path, background_path, output_path, params, is_video_bg)
        )
        thread.daemon = True
        thread.start()

        logger.info(f"任务 {task_id} 开始处理")

    def _process_video_task(self, task_id, human_video_path, background_path,
                            output_path, params, is_video_bg):
        """处理视频任务的后台函数"""
        start_time = time.time()

        try:
            # 创建视频处理器
            processor = self.video_processor.create_processor(task_id)

            # 进度回调函数
            def progress_callback(progress, message):
                self.update_task_progress(task_id, progress, message)

            # 生成结果预览图
            preview_path = output_path.replace('.mp4', '_preview.jpg')
            processor.generate_preview_image(
                human_video_path, background_path, params, preview_path
            )

            # 根据背景类型选择处理方法
            if is_video_bg:
                success, message = processor.process_video_with_video_background(
                    human_video_path, background_path, output_path, params, progress_callback
                )
            else:
                success, message = processor.process_video_with_image_background(
                    human_video_path, background_path, output_path, params, progress_callback
                )

            processing_time = int(time.time() - start_time)

            if success:
                self.complete_task(task_id, output_path, preview_path, processing_time)
                logger.info(f"任务 {task_id} 处理成功，耗时 {processing_time} 秒")
            else:
                self.fail_task(task_id, message, processing_time)
                logger.error(f"任务 {task_id} 处理失败: {message}")

            # 延迟清理处理器
            threading.Timer(300, lambda: self.video_processor.remove_processor(task_id)).start()

        except Exception as e:
            processing_time = int(time.time() - start_time)
            error_msg = f"处理异常: {str(e)}"
            self.fail_task(task_id, error_msg, processing_time)
            logger.error(f"任务 {task_id} 处理异常: {e}")

    def get_active_tasks_count(self):
        """获取活跃任务数量"""
        with self.task_lock:
            return len([t for t in self.tasks.values() if t['status'] == 'processing'])

    def cleanup_completed_tasks(self, hours=24):
        """清理已完成的内存任务（保留数据库记录）"""
        cutoff_time = datetime.now().timestamp() - (hours * 3600)

        with self.task_lock:
            completed_tasks = []
            for task_id, task_info in self.tasks.items():
                if task_info['status'] in ['completed', 'failed']:
                    updated_time = datetime.fromisoformat(task_info['updated_at']).timestamp()
                    if updated_time < cutoff_time:
                        completed_tasks.append(task_id)

            for task_id in completed_tasks:
                self.tasks.pop(task_id, None)
                logger.info(f"从内存中清理已完成任务: {task_id}")

            return len(completed_tasks)


class BatchTaskManager:
    """批量任务管理器"""

    def __init__(self, task_manager: TaskManager):
        self.task_manager = task_manager
        self.batch_tasks = {}  # 批量任务记录
        self.batch_lock = threading.Lock()

    def create_batch_task(self, combinations, params):
        """创建批量任务"""
        batch_id = str(uuid.uuid4())
        task_ids = []

        try:
            for combo in combinations:
                human_file_id = combo.get('human_file_id')
                background_file_id = combo.get('background_file_id')
                output_name = combo.get('output_name', f'batch_output_{len(task_ids)}.mp4')

                # 为每个组合创建单独的任务
                task_id = self.task_manager.create_task(
                    human_file_id, background_file_id, params, output_name
                )

                if task_id:
                    task_ids.append(task_id)
                else:
                    logger.error(f"创建批量任务中的子任务失败: {combo}")

            # 记录批量任务
            with self.batch_lock:
                self.batch_tasks[batch_id] = {
                    'batch_id': batch_id,
                    'task_ids': task_ids,
                    'total_tasks': len(task_ids),
                    'completed_tasks': 0,
                    'failed_tasks': 0,
                    'status': 'created',
                    'created_at': datetime.now().isoformat()
                }

            return batch_id, task_ids

        except Exception as e:
            logger.error(f"创建批量任务失败: {e}")
            return None, []

    def get_batch_status(self, batch_id):
        """获取批量任务状态"""
        with self.batch_lock:
            batch_info = self.batch_tasks.get(batch_id)

            if not batch_info:
                return None

            # 更新子任务状态统计
            completed = 0
            failed = 0
            processing = 0

            for task_id in batch_info['task_ids']:
                task_status = self.task_manager.get_task_status(task_id)
                if task_status:
                    if task_status['status'] == 'completed':
                        completed += 1
                    elif task_status['status'] == 'failed':
                        failed += 1
                    elif task_status['status'] == 'processing':
                        processing += 1

            batch_info.update({
                'completed_tasks': completed,
                'failed_tasks': failed,
                'processing_tasks': processing
            })

            # 更新批量任务状态
            if completed + failed == batch_info['total_tasks']:
                batch_info['status'] = 'completed'
            elif processing > 0:
                batch_info['status'] = 'processing'

            return batch_info

    def start_batch_processing(self, batch_id, upload_dao):
        """开始批量处理"""
        with self.batch_lock:
            batch_info = self.batch_tasks.get(batch_id)

            if not batch_info:
                return False

            batch_info['status'] = 'processing'

        # 为每个子任务启动处理
        for task_id in batch_info['task_ids']:
            task_info = self.task_manager.get_task_status(task_id)
            if not task_info:
                continue

            try:
                # 获取文件路径
                human_upload = upload_dao.get_upload_by_id(task_info['human_file_id'])
                bg_upload = upload_dao.get_upload_by_id(task_info['background_file_id'])

                if not human_upload or not bg_upload:
                    self.task_manager.fail_task(task_id, "文件不存在")
                    continue

                # 生成输出路径
                output_filename = f"{task_id}_{task_info['output_name']}"
                output_path = os.path.join(Config.OUTPUT_FOLDER, output_filename)

                # 开始处理
                self.task_manager.start_processing(
                    task_id, human_upload['file_path'], bg_upload['file_path'],
                    output_path, task_info['params']
                )

            except Exception as e:
                logger.error(f"启动批量任务 {task_id} 失败: {e}")
                self.task_manager.fail_task(task_id, f"启动失败: {str(e)}")

        return True